import asyncio
import argparse
import time
from transformers import AutoTokenizer
# asyncio：Python异步I/O库，用于并发处理请求。
# argparse：解析命令行参数。
# transformers.AutoTokenizer：HuggingFace的tokenizer，用于编码/解码文本，将word和token id互相转化。

import swiftllm

# async def：定义异步函数，允许使用await关键字等待异步操作完成。​这是为了​流式输出​​时，逐个token生成时，异步允许客户端实时接收部分结果
# 这是一个可以在执行过程中“暂停”的函数（比如等待模型生成 Token）。
# 它不会阻塞程序的其他部分，适合处理 I/O 密集型任务（如网络请求、模型推理）。
# 调用时需要用 await（如 await send_request_and_wait_streaming_mode(...)）。
async def send_request_and_wait_streaming_mode(engine: swiftllm.Engine, tokenizer: AutoTokenizer, prompt: str, output_len: int):
    # swiftllm.RawRequest：创建一个原始请求对象，包含输入提示和期望输出长度。
    raw_request = swiftllm.RawRequest(prompt, output_len)
    output_token_ids = [] # 用于存储模型生成的所有 Token ID
    token_latencies = [] # 用于存储每个 Token 的生成延迟时间
    last_time = time.perf_counter() # 记录当前时间（用于计算延迟）

    # engine.add_request_and_stream：SwiftLLM 引擎提供的异步方法，会逐步生成 Token。它返回一个异步生成器，逐个返回生成的 Token。
    # 这里的流式处理允许客户端实时接收部分结果，而不是等待整个请求完成。
    # async for​​：这是一个异步循环，每次迭代会等待下一个 Token 生成
    async for step_output in engine.add_request_and_stream(raw_request):
        output_token_ids.append(step_output.token_id)
        token_latencies.append(time.perf_counter() - last_time)
        last_time = time.perf_counter()

    print("---------------------------------")
    print(f"Prompt: {prompt}")
    print(f"Output: {tokenizer.decode(output_token_ids)}")
    print(f"Token latencies (ms): {[round(t*1000, 1) for t in token_latencies]}")

async def send_request_and_wait_non_streaming_mode(engine: swiftllm.Engine, tokenizer: AutoTokenizer, prompt: str, output_len: int):
    raw_request = swiftllm.RawRequest(prompt, output_len)
    # engine.add_request_and_wait：SwiftLLM 引擎提供的异步方法，会等待整个请求完成并返回所有生成的 Token。
    # await 表示当前协程会暂停，直到 add_request_and_wait 返回结果,在此期间，程序可以处理其他任务（异步并发）
    # 异步是意思是在等待模型生成时，协程不阻塞主线程，CPU可以处理其他任务（如响应其他用户请求）
    # 协程​​：在python中就是可以暂停和恢复的函数，可以在执行到 await 时暂停，并在操作完成后恢复执行
    (_, output_token_ids) = await engine.add_request_and_wait(raw_request)
    print("---------------------------------")
    print(f"Prompt: {prompt}")
    print(f"Output: {tokenizer.decode(output_token_ids)}")

async def main():
    parser = argparse.ArgumentParser()
    parser.description = """
        An example script to demonstrate how to use the swiftllm engine (both streaming and non-streaming mode)
    """
    parser.add_argument(
        "--model-path",
        help="Path to the model. Note: please download the model weights from HuggingFace in advance and specify the path here.",
        type=str,
        required=True
    )
    parser.add_argument(
        "--streaming",
        help="Use streaming mode",
        action="store_true"
    )
    args = parser.parse_args()
    model_path = args.model_path
    is_streaming_mode = args.streaming

    engine_config = swiftllm.EngineConfig(
        model_path = model_path,
        use_dummy = False, # 使用真实模型（非测试 dummy）
        
        block_size = 16, #每个KVCache块包含的token数量
        gpu_mem_utilization = 0.99, #最多允许使用的 GPU 内存比例
        num_cpu_blocks = 1024, # CPU 上预留的备用 KV 缓存块数量（当 GPU 缓存不足时用于溢出存储）
        max_seqs_in_block_table = 128, # 块表中允许的最大序列数（块表用于管理缓存块与序列的映射）
        max_blocks_per_seq = 3072, # 单个序列（请求）允许使用的最大 KV 缓存块数（限制单请求的缓存占用）

        max_batch_size = 4, # 批处理中允许的最大并发序列（请求）数量
        max_tokens_in_batch = 1024 # 批处理中所有序列的总 token 数量上限（防止单批计算量过大）
    )

    prompt_and_output_lens = [
        ("Life blooms like a flower, far away", 10),
        ("one two three four five", 50),
        ("A B C D E F G H I J K L M N O P Q R S T U V", 5),
        ("To be or not to be,", 15),
    ]

    engine = swiftllm.Engine(engine_config) # 创建 SwiftLLM 引擎实例
    tokenizer = AutoTokenizer.from_pretrained(model_path) # 加载tokenizer

    await engine.initialize() # 异步初始化模型（加载权重等）
    asyncio.create_task(engine.start_all_event_loops()) # 启动后台任务处理推理请求（非阻塞主线程）
    # asyncio.create_task() 将每个请求包装为异步任务（立即开始执行）
    tasks = []
    for prompt, output_len in prompt_and_output_lens:
        if is_streaming_mode:
            task = asyncio.create_task(send_request_and_wait_streaming_mode(engine, tokenizer, prompt, output_len))
        else:
            task = asyncio.create_task(send_request_and_wait_non_streaming_mode(engine, tokenizer, prompt, output_len))
        tasks.append(task)
        await asyncio.sleep(0.1) # 间隔0.1秒添加任务，模拟真实请求的时间差
    await asyncio.gather(*tasks) # 等待所有任务完成

if __name__ == "__main__":
    asyncio.run(main())
